#include "Context.hpp"

namespace core
{
    std::unique_ptr<Context> Context::instance_ = nullptr;
    void Context::Init(SDL_Window *window)
    {
        if (!instance_)
        {
            instance_.reset(new Context(window));
        }
    }
    void Context::createVulkanInstance()
    {
        vk::InstanceCreateInfo createInfo;
        vk::ApplicationInfo appInfo;
        appInfo.setPApplicationName("Hello triangle");
        appInfo.setApplicationVersion(VK_MAKE_VERSION(1, 0, 0));
        appInfo.setPEngineName("No engine");
        appInfo.setEngineVersion(VK_MAKE_VERSION(1, 0, 0));
        appInfo.setApiVersion(VK_API_VERSION_1_3);

        createInfo.setPApplicationInfo(&appInfo);
        createInfo.setPpEnabledExtensionNames(getRequiredExtensions(&createInfo.enabledExtensionCount));
        for (int i = 0; i < createInfo.enabledExtensionCount; i++)
        {
            printf("Extension: %s\n", createInfo.ppEnabledExtensionNames[i]);
        }
        vlkInstance = vk::createInstance(createInfo);
        if (!vlkInstance)
        {
            throw std::runtime_error("Create vulkan instance failed!");
        }
    }
    void Context::createSurface(SDL_Window *window)
    {
        if (!SDL_Vulkan_CreateSurface(window, vlkInstance, nullptr, &surface))
        {
            throw std::runtime_error("Create surface failed");
        }
    }
    char const *const *Context::getRequiredExtensions(uint32_t *count)
    {
        /*
         * If validation layers are enabled, we need to request the VK_EXT_DEBUG_UTILS_EXTENSION_NAME extension
         * as well as the extensions required by GLFW.
         *
         * Note: calling this repeatedly could cause a memory leak as we create a new array each time
         * if validation layers are enabled.
         */
        assert(count != NULL);
        // Get the required extensions from SDL and set count to the number of extensions
        char const *const *sdlExtensions = SDL_Vulkan_GetInstanceExtensions(count);
        if (sdlExtensions == NULL)
        {
            throw std::runtime_error("Get SDL instance extensions failed!");
        }
        return sdlExtensions;
    }
    void Context::pickPhycicalDevice()
    {
        if (vlkInstance)
        {
            auto physicalDevices = vlkInstance.enumeratePhysicalDevices();
            vk::PhysicalDevice tempPhyDevice = VK_NULL_HANDLE;
            // phyDevice = physicalDevices[0];
            for (int i = 0; i < physicalDevices.size(); i++)
            {
                // Find queue family
                findQueueFamilyIndicates(physicalDevices[i]);
                std::vector<vk::ExtensionProperties> availableExtensions = physicalDevices[i].enumerateDeviceExtensionProperties();
                bool extensionFound = false;
                for (int j = 0; j < availableExtensions.size(); j++)
                {
                    if (strcmp(VK_KHR_SWAPCHAIN_EXTENSION_NAME, availableExtensions[j].extensionName) == 0)
                    {
                        extensionFound = true;
                        break;
                    }
                }
                if (!extensionFound)
                {
                    continue;
                }
                findSwapChainSupported(physicalDevices[i], surface);
                if (queueFamilyIndicates.hasGraphicFamily && queueFamilyIndicates.hasPresentFamily &&
                    swapChainSupportDetails.formats.size() > 0 && swapChainSupportDetails.presentModes.size() > 0)
                {
                    phyDevice = physicalDevices[i];
                    break;
                }
            }
        }
        if (phyDevice == VK_NULL_HANDLE)
        {
            throw std::runtime_error("Find physical device failed");
        }
        std::cout << "Physical device is " << phyDevice.getProperties().deviceName << std::endl;
    }
    void Context::findQueueFamilyIndicates(vk::PhysicalDevice physicalDevice)
    {
        auto properties = physicalDevice.getQueueFamilyProperties();
        for (int i = 0; i < properties.size(); i++)
        {
            const auto &property = properties[i];
            if (property.queueFlags & vk::QueueFlagBits::eGraphics)
            {
                queueFamilyIndicates.graphicFamily = i;
                queueFamilyIndicates.hasGraphicFamily = true;
            }
            if (physicalDevice.getSurfaceSupportKHR(i, surface))
            {
                queueFamilyIndicates.presentFamily = i;
                queueFamilyIndicates.hasPresentFamily = true;
            }
            if (queueFamilyIndicates.hasGraphicFamily && queueFamilyIndicates.hasPresentFamily)
            {
                break;
            }
        }
    }
    void Context::findSwapChainSupported(vk::PhysicalDevice phyDevice, VkSurfaceKHR surface)
    {
        swapChainSupportDetails.capabilities = phyDevice.getSurfaceCapabilitiesKHR(surface);
        swapChainSupportDetails.formats = phyDevice.getSurfaceFormatsKHR(surface);
        swapChainSupportDetails.presentModes = phyDevice.getSurfacePresentModesKHR(surface);
    }
    void Context::createDevice()
    {
        uint32_t uniqueQueueFamilies[2] = {queueFamilyIndicates.graphicFamily.value(), queueFamilyIndicates.presentFamily.value()};
        uint32_t numOfUniqueQueueFamilies = uniqueQueueFamilies[0] == uniqueQueueFamilies[1] ? 1 : 2;
        std::vector<vk::DeviceQueueCreateInfo> deviceQueueCreateInfos;
        float priority = 1.0;
        for (int i = 0; i < numOfUniqueQueueFamilies; i++)
        {
            vk::DeviceQueueCreateInfo deviceQueueCreateInfo;
            // deviceQueueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
            deviceQueueCreateInfo.setQueueFamilyIndex(uniqueQueueFamilies[i])
                .setQueueCount(1)
                .setQueuePriorities(priority);
            deviceQueueCreateInfos.push_back(deviceQueueCreateInfo);
        }
        vk::DeviceCreateInfo deviceCreateInfo;
        std::array extensions = {VK_KHR_SWAPCHAIN_EXTENSION_NAME};

        deviceCreateInfo.setQueueCreateInfos(deviceQueueCreateInfos)
            .setPEnabledExtensionNames(extensions)
            .setQueueCreateInfoCount(numOfUniqueQueueFamilies)
            .setPEnabledFeatures(nullptr)
            .setEnabledExtensionCount(extensions.size());
        deviceCreateInfo.setEnabledLayerCount(0);
        device = phyDevice.createDevice(deviceCreateInfo);
        if (!device)
        {
            throw std::runtime_error("Create device failed");
        }
        std::cout << "Device is created" << std::endl;
    }
    void Context::getDeviceQueue()
    {
        graphicQueue = device.getQueue(queueFamilyIndicates.graphicFamily.value(), 0);
        presentQueue = device.getQueue(queueFamilyIndicates.presentFamily.value(), 0);
    }
    void Context::createSwapchain(SDL_Window *window)
    {
        if (swapChainSupportDetails.formats.size() == 0)
        {
            throw std::runtime_error("Swap chain support not available (no formats)\n");
        }
        if (swapChainSupportDetails.presentModes.size() == 0)
        {
            throw std::runtime_error("Swap chain support not available (no present modes)\n");
        }
        std::cout << "Swap chain support: "
                  << swapChainSupportDetails.formats.size()
                  << " formats, "
                  << swapChainSupportDetails.presentModes.size()
                  << " present modes\n"
                  << std::endl;
        // Choose the best surface format from the available formats
        vk::SurfaceFormatKHR surfaceFormat = swapChainSupportDetails.formats[0];
        for (int i = 0; i < swapChainSupportDetails.formats.size(); i++)
        {
            if (swapChainSupportDetails.formats[i].format == vk::Format::eB8G8R8A8Srgb && swapChainSupportDetails.formats[i].colorSpace == vk::ColorSpaceKHR::eSrgbNonlinear)
            {
                surfaceFormat = swapChainSupportDetails.formats[i];
                break;
            }
        }
        // Choose the best present mode from the available present modes
        vk::PresentModeKHR presentMode = swapChainSupportDetails.presentModes[0];
        for (int j = 0; j < swapChainSupportDetails.presentModes.size(); j++)
        {
            if (swapChainSupportDetails.presentModes[j] == vk::PresentModeKHR::eMailbox)
            {
                presentMode = swapChainSupportDetails.presentModes[j];
                break;
            }
        }
        swapChainExtent = chooseSwapExtent(window, &swapChainSupportDetails.capabilities);
        std::cout << "Swap chain extent: " << swapChainExtent.width << " x " << swapChainExtent.height << std::endl;

        uint32_t imageCount = swapChainSupportDetails.capabilities.minImageCount + 1;
        if (swapChainSupportDetails.capabilities.maxImageCount > 0 && imageCount > swapChainSupportDetails.capabilities.maxImageCount)
        {
            imageCount = swapChainSupportDetails.capabilities.maxImageCount;
        }

        vk::SwapchainCreateInfoKHR swapChainCreateInfo;
        swapChainCreateInfo.sType = vk::StructureType::eSwapchainCreateInfoKHR;
        swapChainCreateInfo.setSurface(surface)
            .setMinImageCount(imageCount)
            .setImageFormat(surfaceFormat.format)
            .setImageColorSpace(surfaceFormat.colorSpace)
            .setImageExtent(swapChainExtent)
            .setImageArrayLayers(1)
            .setImageUsage(vk::ImageUsageFlagBits::eColorAttachment);
        uint32_t queueFamilyIndices[] = {queueFamilyIndicates.graphicFamily.value(), queueFamilyIndicates.presentFamily.value()};
        if (queueFamilyIndicates.graphicFamily != queueFamilyIndicates.presentFamily)
        {
            swapChainCreateInfo.setImageSharingMode(vk::SharingMode::eConcurrent)
                .setQueueFamilyIndexCount(2)
                .setPQueueFamilyIndices(queueFamilyIndices);
        }
        else
        {
            swapChainCreateInfo.setImageSharingMode(vk::SharingMode::eExclusive);
        }
        swapChainCreateInfo.setPreTransform(swapChainSupportDetails.capabilities.currentTransform)
            .setCompositeAlpha(vk::CompositeAlphaFlagBitsKHR::eOpaque)
            .setPresentMode(presentMode)
            .setClipped(VK_TRUE);
        swapchain = device.createSwapchainKHR(swapChainCreateInfo);
        if (!swapchain)
        {
            throw std::runtime_error("Create swap chain failed");
        }
        swapchainImages = device.getSwapchainImagesKHR(swapchain);
        if (swapchainImages.size() == 0)
        {
            throw std::runtime_error("Get swap chain images failed");
        }
        format = surfaceFormat.format;
        std::cout << "Swap chain created" << std::endl;
    }
    vk::Extent2D Context::chooseSwapExtent(SDL_Window *window, vk::SurfaceCapabilitiesKHR *capabilities)
    {
        if (capabilities->currentExtent.width != 0xFFFFFFFF)
        {
            return capabilities->currentExtent;
        }
        int width, height;
        SDL_GetWindowSize(window, &width, &height);
        vk::Extent2D actualExtent = {
            static_cast<uint32_t>(width),
            static_cast<uint32_t>(height)};
        // Clamp the width and height to the min and max extents
        if (actualExtent.width < capabilities->minImageExtent.width)
        {
            actualExtent.width = capabilities->minImageExtent.width;
        }
        else if (actualExtent.width > capabilities->maxImageExtent.width)
        {
            actualExtent.width = capabilities->maxImageExtent.width;
        }

        if (actualExtent.height < capabilities->minImageExtent.height)
        {
            actualExtent.height = capabilities->minImageExtent.height;
        }
        else if (actualExtent.height > capabilities->maxImageExtent.height)
        {
            actualExtent.height = capabilities->maxImageExtent.height;
        }

        return actualExtent;
    }
    void Context::createImageViews()
    {
        for (int i = 0; i < swapchainImagesCount; i++)
        {
            vk::ImageViewCreateInfo imageViewCreateInfo;
            vk::ComponentMapping mapping;
            vk::ImageSubresourceRange range;
            range.setAspectMask(vk::ImageAspectFlagBits::eColor)
                .setBaseMipLevel(0)
                .setLayerCount(1)
                .setLevelCount(1)
                .setBaseArrayLayer(0);
            imageViewCreateInfo.sType = vk::StructureType::eImageViewCreateInfo;
            imageViewCreateInfo.setImage(swapchainImages.at(i))
                .setViewType(vk::ImageViewType::e2D)
                .setFormat(format)
                .setComponents(mapping)
                .setSubresourceRange(range);
            imageViews[i] = device.createImageView(imageViewCreateInfo);
        }
        std::cout << "Image views created" << std::endl;
    }
    void Context::createRenderPass()
    {
        vk::AttachmentDescription colorAttachment;
        colorAttachment.setFormat(format)
            .setSamples(vk::SampleCountFlagBits::e1)
            .setLoadOp(vk::AttachmentLoadOp::eClear)
            .setStoreOp(vk::AttachmentStoreOp::eStore)
            .setStencilLoadOp(vk::AttachmentLoadOp::eDontCare)
            .setStencilStoreOp(vk::AttachmentStoreOp::eDontCare)
            .setInitialLayout(vk::ImageLayout::eUndefined)
            .setFinalLayout(vk::ImageLayout::ePresentSrcKHR);
        vk::AttachmentReference colorAttachmentRef;
        colorAttachmentRef.setAttachment(0).setLayout(vk::ImageLayout::eColorAttachmentOptimal);
        vk::SubpassDescription subpass;
        subpass.setPipelineBindPoint(vk::PipelineBindPoint::eGraphics)
            .setColorAttachmentCount(1)
            .setColorAttachments(colorAttachmentRef);
        vk::SubpassDependency dependency;
        dependency.setSrcSubpass(VK_SUBPASS_EXTERNAL)
            .setDstSubpass(0)
            .setSrcStageMask(vk::PipelineStageFlagBits::eColorAttachmentOutput)
            .setSrcAccessMask(vk::AccessFlagBits::eNone)
            .setDstStageMask(vk::PipelineStageFlagBits::eColorAttachmentOutput)
            .setDstAccessMask(vk::AccessFlagBits::eColorAttachmentWrite);
        vk::RenderPassCreateInfo renderPassInfo;
        renderPassInfo.sType = vk::StructureType::eRenderPassCreateInfo;
        renderPassInfo.setAttachmentCount(1)
            .setPAttachments(&colorAttachment)
            .setSubpassCount(1)
            .setPSubpasses(&subpass)
            .setDependencyCount(1)
            .setPDependencies(&dependency);
        renderPass = device.createRenderPass(renderPassInfo);
        if (!renderPass)
        {
            throw std::runtime_error("Create render pass failed");
        }
        std::cout << "Render pass created" << std::endl;
    }
    void Context::setUpUniformBufferLayout()
    {
        vk::DescriptorSetLayoutBinding uboLayoutBinding;
        uboLayoutBinding.setBinding(0)
            .setDescriptorType(vk::DescriptorType::eUniformBuffer)
            .setDescriptorCount(1)
            .setStageFlags(vk::ShaderStageFlagBits::eVertex)
            .setPImmutableSamplers(nullptr);
        vk::DescriptorSetLayoutCreateInfo setLayoutCreateInfo;
        setLayoutCreateInfo.sType = vk::StructureType::eDescriptorSetLayoutCreateInfo;
        setLayoutCreateInfo.setBindingCount(1)
            .setPBindings(&uboLayoutBinding);
        descriptorSetLayout = device.createDescriptorSetLayout(setLayoutCreateInfo);
        if (!descriptorSetLayout)
        {
            throw std::runtime_error("Set up uniform buffer layout failed");
        }
    }
    void Context::createGraphicsPipeline()
    {
        size_t vertShaderCodeSize;
        char *vert_shader_code = read_entire_binary_file("shaders/shader.vert.spv", &vertShaderCodeSize);
        printf(" Read %d bytes\n", vert_shader_code_size);

        VkShaderModule vert_shader_module = create_shader_module(vert_shader_code, vert_shader_code_size);
        free(vert_shader_code);

        size_t frag_shader_code_size;
        char *frag_shader_code = read_entire_binary_file("shaders/shader.frag.spv", &frag_shader_code_size);
        printf(" Read %d bytes\n", frag_shader_code_size);

        VkShaderModule frag_shader_module = create_shader_module(frag_shader_code, frag_shader_code_size);
        free(frag_shader_code);

        VkPipelineShaderStageCreateInfo vert_shader_stage_info = {0};
        vert_shader_stage_info.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        vert_shader_stage_info.stage = VK_SHADER_STAGE_VERTEX_BIT;
        vert_shader_stage_info.module = vert_shader_module;
        vert_shader_stage_info.pName = "main";

        VkPipelineShaderStageCreateInfo frag_shader_stage_info = {0};
        frag_shader_stage_info.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        frag_shader_stage_info.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
        frag_shader_stage_info.module = frag_shader_module;
        frag_shader_stage_info.pName = "main";

        VkPipelineShaderStageCreateInfo shader_stages[] = {vert_shader_stage_info, frag_shader_stage_info};

        // Get the vertex format description
        VkVertexInputBindingDescription binding_description = get_binding_description();
        size_t attribute_descriptions_count;
        VkVertexInputAttributeDescription *attribute_descriptions = get_attribute_descriptions(&attribute_descriptions_count);

        VkPipelineVertexInputStateCreateInfo vertex_input_info = {0};
        vertex_input_info.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
        vertex_input_info.vertexBindingDescriptionCount = 1;
        vertex_input_info.vertexAttributeDescriptionCount = attribute_descriptions_count;
        vertex_input_info.pVertexBindingDescriptions = &binding_description;
        vertex_input_info.pVertexAttributeDescriptions = attribute_descriptions;

        VkPipelineInputAssemblyStateCreateInfo input_assembly = {0};
        input_assembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
        input_assembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
        input_assembly.primitiveRestartEnable = VK_FALSE;

        VkPipelineViewportStateCreateInfo viewport_state = {0};
        viewport_state.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
        viewport_state.viewportCount = 1;
        viewport_state.scissorCount = 1;

        VkPipelineRasterizationStateCreateInfo rasterizer = {0};
        rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
        rasterizer.depthClampEnable = VK_FALSE;
        rasterizer.rasterizerDiscardEnable = VK_FALSE;
        rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
        rasterizer.lineWidth = 1.0f;
        rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;
        rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
        rasterizer.depthBiasEnable = VK_FALSE;

        VkPipelineMultisampleStateCreateInfo multisampling = {0};
        multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
        multisampling.sampleShadingEnable = VK_FALSE;
        multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;

        VkPipelineColorBlendAttachmentState color_blend_attachment = {0};
        color_blend_attachment.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
        color_blend_attachment.blendEnable = VK_FALSE;

        VkPipelineColorBlendStateCreateInfo color_blending = {0};
        color_blending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
        color_blending.logicOpEnable = VK_FALSE;
        color_blending.logicOp = VK_LOGIC_OP_COPY;
        color_blending.attachmentCount = 1;
        color_blending.pAttachments = &color_blend_attachment;
        color_blending.blendConstants[0] = 0.0f;
        color_blending.blendConstants[1] = 0.0f;
        color_blending.blendConstants[2] = 0.0f;
        color_blending.blendConstants[3] = 0.0f;

        VkDynamicState dynamic_states[] = {VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR};

        VkPipelineDynamicStateCreateInfo dynamic_state = {0};
        dynamic_state.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
        dynamic_state.dynamicStateCount = 2;
        dynamic_state.pDynamicStates = dynamic_states;

        VkPipelineLayoutCreateInfo pipeline_layout_info = {0};
        pipeline_layout_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
        pipeline_layout_info.setLayoutCount = 1;
        pipeline_layout_info.pSetLayouts = &descriptor_set_layout;
        pipeline_layout_info.pushConstantRangeCount = 0;

        pipeline = device.createPipelineLayout(pipeline_layout_info);
    }



















    vk::Instance &Context::GetVkInstance()
    {
        return vlkInstance;
    }
    VkSurfaceKHR &Context::GetSurface()
    {
        return surface;
    }
    void Context::Quit()
    {
        if (instance_)
        {
            instance_.reset();
        }
    }
    Context::Context(SDL_Window *window)
    {
        /* 1. Create vulkan instance */
        createVulkanInstance();
        /* 2. Create window surface*/
        createSurface(window);
        /* 3. Pickup physical device */
        pickPhycicalDevice();
        /* 4. Create logical device */
        createDevice();
        /* 5. Get device queues */
        getDeviceQueue();
        /* 6. Create the swap chain */
        createSwapchain(window);
        /* 7. Create image views */
        createImageViews();
        /* 8. Create render pass */
        createRenderPass();
        /* 9. Set up uniform buffer layout */
        setUpUniformBufferLayout();
        /* 10. Create the graphics pipeline */
    }
    Context::~Context()
    {
    }
}